import os
import base64
import mimetypes
from typing import Optional, List, Dict
from mcp.server.fastmcp import FastMCP
from zai import ZhipuAiClient


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("ZHIPUAI_API_KEY")
    if not api_key:
        raise ValueError("ZHIPUAI_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


def _convert_local_image_to_base64(image_path: str) -> str:
    """
    将本地图片文件转换为base64格式

    参数:
        image_path: 本地图片文件路径

    返回:
        base64编码的图片数据

    异常:
        ValueError: 文件不存在或无法读取
    """
    if not os.path.exists(image_path):
        raise ValueError(f"本地图片文件不存在: {image_path}")

    try:
        with open(image_path, "rb") as image_file:
            image_data = image_file.read()
            base64_data = base64.b64encode(image_data).decode("utf-8")
            return base64_data
    except Exception as e:
        raise ValueError(f"无法读取本地图片文件 {image_path}: {str(e)}")


def format_multimodal_response(response_data: Dict) -> str:
    """
    格式化多模态响应数据

    参数:
        response_data: 多模态原始响应数据

    返回:
        格式化后的响应字符串
    """
    if "choices" not in response_data or not response_data["choices"]:
        return "API响应格式错误"

    message = response_data["choices"][0].get("message", {})
    content = message.get("content", "")

    return content if content else "响应内容为空"


def _classify_file_type(url: str) -> str:
    """
    根据文件URL或路径判断文件类型

    参数:
        url: 文件URL或本地路径

    返回:
        文件类型: 'image', 'video', 'file', 'text'
    """
    # 提取文件后缀
    if url.startswith(("http://", "https://")):
        # URL格式，从路径中提取后缀
        path_part = url.split("?")[0]  # 移除查询参数
        extension = os.path.splitext(path_part)[1].lower()
    else:
        # 本地文件路径
        extension = os.path.splitext(url)[1].lower()

    # 图片格式
    image_extensions = {
        ".jpg",
        ".jpeg",
        ".png",
        ".gif",
        ".bmp",
        ".webp",
        ".tiff",
        ".svg",
    }
    # 视频格式
    video_extensions = {".mp4", ".avi", ".mov", ".wmv", ".flv", ".mkv", ".webm", ".m4v"}
    # 文件格式
    file_extensions = {
        ".pdf",
        ".txt",
        ".doc",
        ".docx",
        ".xls",
        ".xlsx",
        ".ppt",
        ".pptx",
        ".csv",
        ".json",
        ".xml",
        ".html",
        ".htm",
        ".md",
        ".rtf",
    }

    if extension in image_extensions:
        return "image"
    elif extension in video_extensions:
        return "video"
    elif extension in file_extensions:
        return "file"


mcp = FastMCP("ZhipuAI Multimodal MCP Server")


@mcp.tool()
async def understand_multi_media(
    urls: List[str], question: str = "请分析这些多媒体内容"
) -> str:
    """
    多媒体内容理解（使用智谱AI glm-4.5v模型）
    支持图片、视频、文件和文本的智能分类处理

    参数:
        urls: 多媒体URL列表或本地文件路径列表，例如：
              ["https://example.com/image1.jpg", "/path/to/local/video.mp4", "https://example.com/document.pdf"]
        question: 分析问题，默认值："请分析这些多媒体内容"

    返回:
        多媒体理解结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        Exception: API调用失败
    """
    if not urls:
        raise ValueError("必须提供至少一个多媒体URL")

    api_key = get_api_key()
    client = ZhipuAiClient(api_key=api_key)

    # 构建内容数组
    content = []
    video_count = 0

    if urls is None or not isinstance(urls, list):
        # 说明是text
        pass
    else:
        for url in urls:
            file_type = _classify_file_type(url)

            if file_type == "image":
                # 处理图片
                if url.startswith(("http://", "https://")):
                    # HTTP URL，直接使用
                    content.append({"type": "image_url", "image_url": {"url": url}})
                else:
                    # 本地文件路径，转换为base64
                    try:
                        base64_data = _convert_local_image_to_base64(url)
                        content.append(
                            {"type": "image_url", "image_url": {"url": base64_data}}
                        )
                    except ValueError as e:
                        raise ValueError(f"本地图片处理失败: {str(e)}")

            elif file_type == "video":
                # 处理视频
                video_count += 1
                if video_count > 1:
                    raise ValueError("视频只能有一个")
                if not url.startswith(("http://", "https://")):
                    raise ValueError("视频只支持URL格式，不支持本地文件")
                content.append({"type": "video_url", "video_url": {"url": url}})

            elif file_type == "file":
                # 处理文件
                if not url.startswith(("http://", "https://")):
                    raise ValueError("文件只支持URL格式，不支持本地文件")
                content.append({"type": "file_url", "file_url": {"url": url}})

    # 只有在没有文本内容时才添加question
    content.append({"type": "text", "text": question})

    try:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=[{"role": "user", "content": content}],
            thinking={"type": "enabled"},
        )

        # 将SDK响应转换为字典格式
        response_dict = {
            "choices": [{"message": {"content": response.choices[0].message.content}}]
        }

        return format_multimodal_response(response_dict)
    except Exception as e:
        raise Exception(f"多媒体理解失败: {str(e)}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
